#! /usr/bin/env python3
###########################################################################
# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: MIT
###########################################################################
import argparse
import os
import os.path
import re
import selectors
import socket
import struct
import sys
import time

script_dir = os.path.realpath(os.path.dirname(__file__))
top_dir = os.path.dirname(script_dir)
sys.path.append(os.path.join(top_dir, 'py-lib'))

from qemu_socket import QemuSocket, LISTEN, CONNECT
from utils import dump_packet, parse_time

_raw_address_re = re.compile(r'raw:([^:]+)(?::([^:]+))?\Z')
_qemu_address_re = re.compile(r'(?:qemu:)?([^:]*):(\d+)\Z')

_int_re = re.compile(r'[0-9]+\Z')
_byte_array_re = re.compile(r'[0-9a-f]{1,2}(:[0-9a-f]{1,2})*\Z')

###########################################################################

SB_CONTROL_PORT = 0xfe
ETHERTYPE_NANOTUBE_CONTROL = 0xf00d

class CapsuleError(RuntimeError):
  __slots__ = []

class HeaderFormat:
  __slots__ = []

  def get_field(self, params, name, default=None):
    val = params.get(name, default)
    if val is None:
      raise CapsuleError("Unspecified '%s' parameter." % name)
    return val

  def pack(self, params):
    vals = []
    for f_name, f_default in self.field_defs:
      val = self.get_field(params, f_name, f_default)
      vals.append(val)

    return struct.pack(self.format, *vals)

  def unpack(self, params, buf, offset):
    hdr_len = struct.calcsize(self.format)
    hdr_end = offset+hdr_len
    if len(buf) < hdr_end:
      raise CapsuleError("Capsule is too short for a %s." % self.name)

    view = memoryview(buf)[offset:hdr_end]
    vals = struct.unpack(self.format, view)
    field_names = [n for n,d in self.field_defs]
    fields = list(zip(field_names, vals))
    return (fields, hdr_len)

class SbCapsuleFormat(HeaderFormat):
  __slots__ = []
  name = 'SbCapsule'
  format = '<B'
  field_defs = [
    ('port', SB_CONTROL_PORT),
  ]

  def unpack(self, params, buf, offset):
    res = super().unpack(params, buf, offset)
    if res == None:
      return None

    fields, hdr_len = res
    if dict(fields)['port'] != SB_CONTROL_PORT:
      return None

    return res

class EthEncapFormat(HeaderFormat):
  __slots__ = []
  name = 'EthEncap'
  format = '!6s6sH'
  field_defs = [
    ('dhost',      b'\xff\xff\xff\xff\xff\xff'),
    ('shost',      b'\x00\x00\x00\x00\x00\x00'),
    ('ethertype',  ETHERTYPE_NANOTUBE_CONTROL),
  ]

  def unpack(self, params, buf, offset):
    res = super().unpack(params, buf, offset)
    if res == None:
      return None

    fields, hdr_len = res
    if dict(fields)['ethertype'] != ETHERTYPE_NANOTUBE_CONTROL:
      return None

    return res

class NanotubeControlFormat(HeaderFormat):
  __slots__ = []
  name = 'NanotubeControl'
  format = '<HHH'
  field_defs = [
    ('request_id', 0),
    ('resource_id', None),
    ('response_code', 1),
  ]

OPCODE_READ=0
OPCODE_INSERT=1
OPCODE_UPDATE=2
OPCODE_WRITE=3
OPCODE_REMOVE=4

class MapControlFormat(HeaderFormat):
  __slots__ = []
  name = 'MapControl'
  format = '<H'
  field_defs = [
    ('opcode', None),
  ]

  def pack(self, params):
    key = self.get_field(params, 'key')
    value = self.get_field(params, 'value')
    buf = bytearray(super().pack(params))
    buf.extend(key)
    buf.extend(value)
    return buf

  def unpack(self, params, buf, offset):
    key = self.get_field(params, 'key')
    value = self.get_field(params, 'value')

    fields, prefix_len = super().unpack(params, buf, offset)
    offset += prefix_len
    view = memoryview(buf)[offset:]

    key_len = len(key)
    value_len = len(value)
    data_len = key_len + value_len
    if len(view) < data_len:
      raise CapsuleError("Capsule is too short for %s.  Have %d but"
                         " need %d." % (self.name, len(view), data_len))

    fields.append(('key', bytes(view[:key_len])))
    fields.append(('value', bytes(view[key_len:data_len])))
    return (fields, prefix_len + data_len)

encap_class_dict = {
  'sb': SbCapsuleFormat,
  'eth': EthEncapFormat,
}

class CapsuleFormat:
  def __init__(self, encap_format):
    self.__layers = [ encap_format,
                      NanotubeControlFormat(),
                      MapControlFormat() ]

  def pack(self, params):
    buf = bytearray()
    for layer in self.__layers:
      buf.extend(layer.pack(params))
    return buf

  def unpack(self, params, buf, offset):
    all_fields = []
    total_len = 0
    for layer in self.__layers:
      res = layer.unpack(params, buf, offset)
      if res == None:
        return None
      l_fields, l_hdr_len = res
      all_fields.extend(l_fields)
      offset += l_hdr_len
      total_len += l_hdr_len

    return (all_fields, total_len)

###########################################################################

class RawSocket:
  def __init__(self, dev_name, rx_enable):
    ether_type = 0
    if rx_enable:
      ether_type = 3
    self.__socket = socket.socket(socket.AF_PACKET, socket.SOCK_RAW)
    self.__socket.bind((dev_name, ether_type))
    self.__socket.setblocking(False)
    self.__packet = None

  def fileno(self):
    return self.__socket.fileno()

  def send(self, packet=None):
    """Send some packet data.

    Either send a new packet or continue sending a previous packet.
    The first call to this method should pass a packet as an argument.
    Subsequent calls for the same packet should not provide an
    argument.  The method will return a flag indicating whether it is
    ready to accept the next packet.  If it returns false then the
    caller should keep calling the method with no argument when the
    socket is writable.  If it returns true then the caller may call
    the method again with another packet to send.
    """

    # The underlying send operation will not send a partial packet, so
    # we just need to identify the packet to send and make sure
    # there's exactly one.
    if packet is not None:
      # Send a new packet
      assert self.__packet is None
    else:
      # Continue sending the previous packet.
      assert self.__packet is not None
      packet = self.__packet
      self.__packet = None

    # Try to send the packet.  Store for later if it fails.
    try:
      self.__socket.send(packet)
    except socket.timeout:
      self.__packet = packet
      return False

    return True

  def receive(self):
    """Receive some packet data.

    Called when data can be read from the socket.  Returns an
    interator over received packets.  Each packet is represented as a
    bytearray.  If the iterator yields None then EOF has been reached
    and the receive method should not be called again.
    """

    while True:
      try:
        res = self.__socket.recvmsg(65536)
        pkt,anc,flags,addr = res
        if addr[2] != socket.PACKET_OUTGOING:
          yield bytearray(pkt)
      except BlockingIOError:
        return

###########################################################################

class ParseError(RuntimeError):
  __slots__ = []

class Actions:
  def __init__(self, capsule_format):
    self.__capsule_format = capsule_format
    self.__params = {}
    self.__send_capsules = []
    self.__expect_params = []

  def get_num_capsules(self):
    return len(self.__send_capsules)

  def pop_send_capsule(self):
    buf,params = self.__send_capsules.pop(0)
    self.__expect_params.append(params)
    return buf

  def peek_expect_params(self):
    return self.__expect_params[0]

  def pop_expect_params(self):
    return self.__expect_params.pop(0)

  def done(self):
    return (len(self.__send_capsules) == 0 and
            len(self.__expect_params) == 0)

  def parse_int(self, name, val):
    m = _int_re.match(val)
    if m is None:
      raise ParseError("Invalid %s %r." % (name, val))

    return int(val)

  def parse_byte_array(self, name, val):
    m = _byte_array_re.match(val)
    if m is None:
      raise ParseError("Invalid %s %r." % (name, val))

    return bytearray(int(x, 16) for x in val.split(":"))

  def set_request_id(self, val):
    self.__params['request_id'] = self.parse_int("request ID", val)

  def set_map_id(self, val):
    self.set_resource_id(val)

  def set_resource_id(self, val):
    self.__params['resource_id'] = self.parse_int("resource ID", val)

  def set_key(self, val):
    self.__params['key'] = self.parse_byte_array("key", val)

  def set_key_len(self, val):
    v = self.parse_int("key length", val)
    self.__params['key'] = bytearray(0 for i in range(v))

  def set_value(self, val):
    self.__params['value'] = self.parse_byte_array("value", val)

  def set_val_len(self, val):
    v = self.parse_int("value length", val)
    self.__params['value'] = bytearray(0 for i in range(v))

  def do_read(self):
    params = dict(self.__params)
    params['opcode'] = OPCODE_READ
    buf = self.__capsule_format.pack(params)
    self.__send_capsules.append((buf, params))

  def do_insert(self):
    params = dict(self.__params)
    params['opcode'] = OPCODE_INSERT
    buf = self.__capsule_format.pack(params)
    self.__send_capsules.append((buf, params))

  def do_update(self):
    params = dict(self.__params)
    params['opcode'] = OPCODE_UPDATE
    buf = self.__capsule_format.pack(params)
    self.__send_capsules.append((buf, params))

  def do_write(self):
    params = dict(self.__params)
    params['opcode'] = OPCODE_WRITE
    buf = self.__capsule_format.pack(params)
    self.__send_capsules.append((buf, params))

  def do_remove(self):
    params = dict(self.__params)
    params['opcode'] = OPCODE_REMOVE
    buf = self.__capsule_format.pack(params)
    self.__send_capsules.append((buf, params))

###########################################################################

# Usage:
#   map_access :12345 map_id=3 key=1:bf:3:cd val_len=8 read
#   map_access :12345 map_id=3 key=4:3:2:1 value=9:8:7:6:5:4:3:2 write

class app:
  def __init__(self, argv):
    p = argparse.ArgumentParser()

    p.add_argument("-e", "--encap",
                   action="store", default="sb",
                   help="Set the encap format.")
    p.add_argument("-m", "--mtu",
                   action="store", type=int, default=1514,
                   help="Set the Maximum Transmit Unit.")
    p.add_argument("-t", "--timeout",
                   action="store", default="1s",
                   help="Set the timeout (s/ms/us/ns).")
    p.add_argument("-w", "--wait-time",
                   action="store", default="1ms",
                   help="Set the extra time to wait (s/ms/us/ns).")
    p.add_argument("-v", "--verbose",
                   action="count", default=0,
                   help="Produce more output.")

    p.add_argument("address", nargs=1,
                   help="The address to connect to.")
    p.add_argument("request", nargs='+',
                   help="The request to send.")

    self.__prog = argv[0]
    self.__args = p.parse_args(argv[1:])

    self.__timeout = parse_time(self.__args.timeout)
    self.__verbosity = self.__args.verbose
    self.__wait_time = parse_time(self.__args.wait_time)

    self.__sel = selectors.DefaultSelector()
    self.__tx_socket = None
    self.__rx_socket = None
    self.__done = False
    self.__fail = False
    self.__time_now = None
    self.__end_time = None

    encap_class = encap_class_dict.get(self.__args.encap)
    if encap_class is None:
      sys.stderr.write("%s: Unknown encap format %r.\n" %
                       (self.__prog, self.__args.encap))
      sys.exit(1)

    self.__capsule_format = CapsuleFormat(encap_class())

    self.__actions = Actions(self.__capsule_format)

    self.parse_address()

    self.process_args()

  def parse_address(self):
    addr_str = self.__args.address[0]
    m = _raw_address_re.match(addr_str)
    if m is not None:
      self.__socket_mode = "raw"
      self.__tx_netdev = m.group(1)
      self.__rx_netdev = m.group(2)
      return

    m = _qemu_address_re.match(addr_str)
    if m is None:
      sys.stderr.write("%s: Invalid address %r.\n" %
                       (self.__prog, addr_str))
      sys.exit(1)

    self.__socket_mode = "qemu"
    self.__hostname = m.group(1)
    self.__port = self.__actions.parse_int("port", m.group(2))

  def process_args(self):
    failed=False
    for arg in self.__args.request:
      s = arg.split("=", 1)
      if len(s) == 2:
        var,val = s

        # Use a method like set_map, set_key etc...
        method = getattr(self.__actions, "set_"+var, None)
        if method is None:
          sys.stderr.write("%s: Unrecognized parameter %r.\n" %
                           (self.__prog, var))
          failed = True
        else:
          try:
            method(val)
          except ParseError as e:
            sys.stderr.write("%s: %s\n" % (self.__prog, str(e)))
            failed = True
        continue

      # Use a method like do_read, do_write etc...
      method = getattr(self.__actions, "do_"+arg, None)
      if method is None:
        sys.stderr.write("%s: Unrecognized action %r.\n" %
                         (self.__prog, arg))
        failed = True
      else:
        try:
          method()
        except RuntimeError as e:
          sys.stderr.write("%s: %s\n" % (self.__prog, str(e)))
          failed = True

    if failed:
      sys.exit(1)

  def show_fields(self, fields, prefix):
    for name, val in fields:
      val_str = repr(val)
      if isinstance(val, (bytearray, bytes)):
        val_str = "%r (%s)" % (val, ":".join("%02x"%x for x in val))
      print("%s%-16s %s" % (prefix, name+":", val_str))

  def get_time(self):
    return time.clock_gettime(time.CLOCK_MONOTONIC)

  def connect_qemu_socket(self):
    sk = QemuSocket(hostname=self.__hostname,
                    port=self.__port,
                    mode=CONNECT,
                    mtu=self.__args.mtu,
                    timeout=self.__args.timeout,
                    verbosity=self.__verbosity-1)
    self.__tx_socket = sk

    while True:
      self.__time_now = self.get_time()
      timeout = self.__end_time - self.__time_now
      if sk.try_connect():
        return

      if timeout < 0:
        sys.stderr.write("%s: Timed out connecting to %s.\n" %
                         (self.__prog, self.__args.address[0]))
        sys.exit(1)

      elapsed = self.__time_now - self.__start_time
      if elapsed > 0:
        time.sleep(min(elapsed/4, 1))

  def connect_raw_socket(self):
    if self.__rx_netdev is not None:
      self.__tx_socket = RawSocket(self.__tx_netdev, False)
      self.__rx_socket = RawSocket(self.__rx_netdev, True)
    else:
      self.__tx_socket = RawSocket(self.__tx_netdev, True)

  def connect_socket(self):
    if self.__verbosity >= 1:
      sys.stderr.write("Connecting to %s.\n" %
                       (self.__args.address[0],))

    self.__start_time = self.get_time()
    self.__end_time = self.__start_time + self.__timeout

    if self.__socket_mode == "qemu":
      self.connect_qemu_socket()
    else:
      assert self.__socket_mode == "raw"
      self.connect_raw_socket()

    if self.__rx_socket is None:
      events = selectors.EVENT_READ | selectors.EVENT_WRITE
      self.__sel.register(self.__tx_socket, events)
      self.__rx_socket = self.__tx_socket
    else:
      events = selectors.EVENT_WRITE
      self.__sel.register(self.__tx_socket, events)
      events = selectors.EVENT_READ
      self.__sel.register(self.__rx_socket, events)

  def disable_tx(self):
    if self.__tx_socket is None:
      return

    key = self.__sel.get_key(self.__tx_socket)
    events = key.events & ~selectors.EVENT_WRITE
    if events == 0:
      self.__sel.unregister(self.__tx_socket)
      self.__tx_socket = None
    else:
      self.__sel.modify(self.__tx_socket, events)

  def send_capsules(self):
    if self.__verbosity >= 1:
      sys.stderr.write("Sending %d capsules.\n" %
                       self.__actions.get_num_capsules())

    while self.__actions.get_num_capsules() != 0:
      p = self.__actions.pop_send_capsule()
      if not self.__tx_socket.send(p):
        return

    if self.__verbosity >= 1:
      sys.stderr.write("No more capsules to send.\n")

    self.disable_tx()

  def handle_read_event(self):
    if self.__verbosity >= 1:
      sys.stderr.write("Receiving capsules.\n")
    for pkt in self.__rx_socket.receive():
      if pkt is None:
        sys.stderr.write("%s: Unexpected EOF from the socket.\n" %
                         self.__prog)
        self.__done = True
        self.__fail = True
        return

      params = self.__actions.peek_expect_params()
      if params is None:
        sys.stderr.write("%s: Received unexpected capsule.\n" %
                         self.__prog)
        dump_packet(sys.stderr, "> ", pkt)
        sys.exit(1)

      print("Received response:")
      res = self.__capsule_format.unpack(params, pkt, 0)
      if res == None:
        return

      fields, tot_len = res
      for name,val in fields:
        print("  %-16s %s" % (name+":", val))

      self.__actions.pop_expect_params()

  def handle_write_event(self):
    if self.__verbosity >= 1:
      sys.stderr.write("Continuing send operation.\n")
    if self.__tx_socket.send():
      self.send_capsules()

  def recv_capsules(self):
    while not self.__done:
      self.__time_now = self.get_time()
      timeout = self.__end_time - self.__time_now
      results = self.__sel.select(max(timeout, 0))
      for key,events in results:
        if (events & selectors.EVENT_READ) != 0:
          assert key.fileobj is self.__rx_socket
          self.handle_read_event()
        if (events & selectors.EVENT_WRITE) != 0:
          assert key.fileobj is self.__tx_socket
          self.handle_write_event()

      if self.__actions.done():
        return

      if timeout < 0:
        if self.__actions.get_num_capsules() == 0:
          action = "receiving"
        else:
          action = "sending"
        sys.stderr.write("%s: Timed out %s capsules.\n" %
                         (self.__prog, action))
        sys.exit(1)

  def extra_wait(self):
    self.__time_now = self.get_time()
    self.__end_time = self.__time_now + self.__wait_time

    self.disable_tx()

    while not self.__done:
      self.__time_now = self.get_time()
      timeout = self.__end_time - self.__time_now
      results = self.__sel.select(max(timeout, 0))
      for key,events in results:
        assert key.fileobj is self.__rx_socket
        if (events & selectors.EVENT_READ) != 0:
          self.handle_read_event()

      if timeout < 0:
        return

  def report_results(self):
    if self.__fail:
      sys.stderr.write("%s: Operations failed.\n" % self.__prog)
      sys.exit(1)

  def run(self):
    self.connect_socket()
    self.send_capsules()
    self.recv_capsules()
    self.extra_wait()
    self.report_results()

app(sys.argv).run()
